#include <WinSock2.h>
#include <iphlpapi.h>
#include "capture.h"

#include "stdio.h"

#pragma comment(lib, "Ws2_32.lib")
#pragma comment(lib, "IPHLPAPI.lib")

#define SERVICE_PORT 56789
#define COMMOND_PORT 45678

void get_addr_and_mask(DWORD *addr, DWORD *mask)
{
#if WINVER >= 0x0600
	sockaddr_in unicastAddr;
	PIP_ADAPTER_ADDRESSES adapterAddresses = NULL;
	ULONG outBufferLength = 0;
	if (GetAdaptersAddresses(AF_INET, 0, NULL, NULL, &outBufferLength) != ERROR_BUFFER_OVERFLOW)
		return;
	adapterAddresses = (PIP_ADAPTER_ADDRESSES)malloc(outBufferLength);
	if (GetAdaptersAddresses(AF_INET, 0, NULL, adapterAddresses, &outBufferLength) != NO_ERROR)
		return;
	for (; adapterAddresses; adapterAddresses = adapterAddresses->Next) {
		if (adapterAddresses->FirstUnicastAddress != NULL) {
			unicastAddr = *(sockaddr_in *)adapterAddresses->FirstUnicastAddress->Address.lpSockaddr;
			*addr = unicastAddr.sin_addr.s_addr;
			*mask = 0;
			while (adapterAddresses->FirstUnicastAddress->OnLinkPrefixLength--)
				*mask |= (1 << (31 - adapterAddresses->FirstUnicastAddress->OnLinkPrefixLength));
			*mask = htonl(*mask);
			break;
		}
	}
#else
	PIP_ADAPTER_INFO adapterInfos;
	ULONG outBufferLength = 0;
	if (GetAdaptersInfo(NULL, &outBufferLength) != ERROR_BUFFER_OVERFLOW)
		return;
	adapterInfos = (PIP_ADAPTER_INFO)malloc(outBufferLength);
	if (GetAdaptersInfo(adapterInfos, &outBufferLength) != NO_ERROR)
		return;
	for (; adapterInfos; adapterInfos = adapterInfos->Next) {
		*addr = inet_addr(adapterInfos->IpAddressList.IpAddress.String);
		*mask = inet_addr(adapterInfos->IpAddressList.IpMask.String);
		break;
	}
#endif
}

DWORD WINAPI broadcast_server(LPVOID lpParameter)
{
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2,2), &wsaData);

	DWORD addr, mask, broadcast_addr;
	get_addr_and_mask(&addr, &mask);
	broadcast_addr = addr | (~mask);

	SOCKET updsocket;
	struct sockaddr_in udpserver;

	udpserver.sin_addr.s_addr = broadcast_addr;
	udpserver.sin_family = AF_INET;
	udpserver.sin_port = ntohs(SERVICE_PORT);

	DWORD port;
	port = ntohl(SERVICE_PORT);

	updsocket = socket(AF_INET, SOCK_DGRAM, IPPROTO_UDP);
	if (updsocket == INVALID_SOCKET)
		WSACleanup();

	char send_buffer[8];

	memcpy(send_buffer, &addr, sizeof(addr));
	memcpy(send_buffer+sizeof(addr), &port, sizeof(port));

	for (;;) {
		sendto(updsocket, send_buffer, sizeof(addr)+sizeof(port), 0, (sockaddr *)&udpserver, sizeof(udpserver));
		Sleep(1000);
	}

	return 0;
}

DWORD WINAPI recv_client_para(LPVOID lpParameter)
{
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2,2), &wsaData);

	SOCKET nativeSocket, remoteSocket;
	struct sockaddr_in service;

	nativeSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (nativeSocket == INVALID_SOCKET) {
		WSACleanup();
		return -1;
	}

	service.sin_addr.s_addr = inet_addr("0.0.0.0");
	service.sin_family = AF_INET;
	service.sin_port = ntohs(COMMOND_PORT);

	if (bind(nativeSocket, (SOCKADDR*)&service, sizeof(service)) == SOCKET_ERROR) {
		closesocket(nativeSocket);
		WSACleanup();
		return -1;
	}

	if (listen(nativeSocket, SOMAXCONN) == SOCKET_ERROR) {
		closesocket(nativeSocket);
		WSACleanup();
		return -1;
	}

	for (;;) {
		struct sockaddr_in remoteaddr;
		remoteSocket = accept(nativeSocket, (struct sockaddr *)&remoteaddr, NULL);
		if (remoteSocket == INVALID_SOCKET) {
			closesocket(nativeSocket);
			WSACleanup();
			return -1;
		}

		for (;;) {
			char recv_buffer[12];
			int ret = recv(remoteSocket, recv_buffer, 4, 0);
			if (ret == 0 || ret == SOCKET_ERROR) {
				break;
			}

			short type = ntohs(*(short *)recv_buffer);
			short size = ntohs(*(short *)(recv_buffer + 2));
			ret = recv(remoteSocket, recv_buffer, size, 0);
			if (ret == 0 || ret == SOCKET_ERROR) {
				break;
			}

			ret = send(remoteSocket, recv_buffer, 2, 0);
			if (ret == SOCKET_ERROR) {
				break;
			}

			if (type == 0x01) {
				int width = ntohl(*(int *)recv_buffer);
				int height = ntohl(*(int *)(recv_buffer + 4));

				set_remote_size(width, height);
			}
			else if (type == 0x02) {
				int x = ntohl(*(int *)recv_buffer);
				int y = ntohl(*(int *)(recv_buffer + 4));

				set_remote_pos(x, y);
			}
			else if (type == 0x03) {
				int x = ntohl(*(int *)recv_buffer);
				int y = ntohl(*(int *)(recv_buffer + 4));
				SetCursorPos(gPosX + x, gPosY + y);
				mouse_event(MOUSEEVENTF_LEFTDOWN, 0, 0, 0, 0);
			}
			else if (type == 0x04) {
				int x = ntohl(*(int *)recv_buffer);
				int y = ntohl(*(int *)(recv_buffer + 4));
				SetCursorPos(gPosX + x, gPosY + y);
				mouse_event(MOUSEEVENTF_LEFTUP, 0, 0, 0, 0);
			}
			else if (type == 0x05) {
				int x = ntohl(*(int *)recv_buffer);
				int y = ntohl(*(int *)(recv_buffer + 4));
				SetCursorPos(gPosX + x, gPosY + y);
				mouse_event(MOUSEEVENTF_LEFTDOWN, 0, 0, 0, 0);
				Sleep(10);
				mouse_event(MOUSEEVENTF_LEFTUP, 0, 0, 0, 0);
			}
			else if (type == 0x06) {
				int x = ntohl(*(int *)recv_buffer);
				int y = ntohl(*(int *)(recv_buffer + 4));
				SetCursorPos(gPosX + x, gPosY + y);
			}
			else if (type == 0x07) {
				int code = ntohl(*(int *)recv_buffer);
				keybd_event(code, 0, 0, 0);
				Sleep(10);
				keybd_event(code, 0, KEYEVENTF_KEYUP, 0);
			}
		}
	}

	closesocket(nativeSocket);
	WSACleanup();

	return 0;
}

DWORD WINAPI send_server_data(LPVOID lpParameter)
{
	WSADATA wsaData;
	WSAStartup(MAKEWORD(2,2), &wsaData);

	SOCKET nativeSocket, remoteSocket;
	struct sockaddr_in service;

	nativeSocket = socket(AF_INET, SOCK_STREAM, IPPROTO_TCP);
	if (nativeSocket == INVALID_SOCKET) {
		WSACleanup();
		return -1;
	}

	service.sin_addr.s_addr = inet_addr("0.0.0.0");
	service.sin_family = AF_INET;
	service.sin_port = ntohs(SERVICE_PORT);

	if (bind(nativeSocket, (SOCKADDR*)&service, sizeof(service)) == SOCKET_ERROR) {
		closesocket(nativeSocket);
		WSACleanup();
		return -1;
	}

	if (listen(nativeSocket, SOMAXCONN) == SOCKET_ERROR) {
		closesocket(nativeSocket);
		WSACleanup();
		return -1;
	}

	//////////////////////////////////////////////////////////////////////////
	BYTE *pImage, *pCompressedImage;
	pImage = NULL;
	pCompressedImage = NULL;
	init_capture();
	set_remote_size(100, 100);
	reset_remote_size(&pImage, &pCompressedImage);

	DWORD tick0, tick1, allTick0, allTick1;
	allTick0 = tick0 = GetTickCount();

	for (;;) {
		printf("listening...\n");

		struct sockaddr_in remoteAddr;
		remoteSocket = accept(nativeSocket, (struct sockaddr *)&remoteAddr, NULL);
		if (remoteSocket == INVALID_SOCKET) {
			closesocket(nativeSocket);
			WSACleanup();
			return -1;
		}

		for (;;) {
			if (has_remote_reset() == TRUE) {
				clean_remote_data(pImage, pCompressedImage);
				reset_remote_size(&pImage, &pCompressedImage);
			}

			//////////////////////////////////////////////////////////////////////////
			capture_screen(pImage);
			tick1 = GetTickCount();
			//printf("capture cost: %d\n", tick1 - tick0);
			tick0 = tick1;

			//////////////////////////////////////////////////////////////////////////
			int size = compress_buffer(pImage, pCompressedImage+4);
			*(int *)pCompressedImage = htonl(size);
			tick1 = GetTickCount();
			//printf("compress cost: %d\n", tick1 - tick0);
			tick0 = tick1;

			int ret = send(remoteSocket, (char *)pCompressedImage, size+4, 0);
			if (ret == 0 || ret == SOCKET_ERROR) {
				break;
			}
			tick1 = GetTickCount();
			//printf("send cost: %d\n", tick1 - tick0);
			tick0 = tick1;

			int recvSize;
			ret = recv(remoteSocket, (char *)&recvSize, 4, 0);
			if (ret == 0 || ret == SOCKET_ERROR) {
				break;
			}
			if (ntohl(recvSize) != size) {
				printf("something error?\n");
			}
			tick1 = GetTickCount();
			//printf("recv cost: %d\n", tick1 - tick0);
			tick0 = tick1;

			allTick1 = GetTickCount();
			printf("all cost: %d\n", allTick1 - allTick0);
			allTick0 = allTick1;

			//printf("%d, %d, %.2f%%\n", gRemoteSize, size, (double)size*100/gRemoteSize);
		}
	}

	closesocket(nativeSocket);
	WSACleanup();

	return 0;
}
